# Date: 2021/4/26
# Author: zuoyuwei
# Purpose: 用于创建四类标签图像数据，包括可见行，不可见行，可见列，不可见列
# Usage: python create_gt.py F:/laibo/Data/tijianbaogao/raw/z_hanglie/img_json/ F:/laibo/Data/tijianbaogao/train_data/
# Modify date: 2021/6/8
# Usage: 将训练过程中的数据处理全部拿到外面来做
# Modify date: 2021/6/16
# Usage: 处理训练过程中的图像数据像素值变化问题->将jpg改为png
# Modify date: 2021/7/14
# Usage：cv2.resize时若采用INTER_AREA方法会产生0-255各个像素值，改为INTER_NEAREST
# Modify date: 2021/7/19
# Usage：将图像和标签的尺寸都扩大一倍，标签改为512*512大小，运行时内存受限
# Modify date: 2021/10/27
# Usage: 随机对生成的图像数据添加数据集扩充操作(数据增强尽量最大程度保留原始图像的色彩、形状和尺寸)
# Modify date: 2021/11/16
# Usage: 插值方法改为INTER_AREA,重新生成标签,直接生成用于训练的256*256大小
# Modify date: 2021/11/24
# Usage: 修改旋转增强图像的标签生成方式（仅旋转每个四边形的四个顶点）
import os
import cv2
import sys
import glob
import json
import math
import numpy as np
from PIL import Image

def rotate_bound(image, angle):
    '''
    旋转后保证图像尺寸变大，保证图像都可以出现
    :param image: 原图
    :param angle: 角度
    :return: 旋转后的图像
    '''
    (h, w) = image.shape[:2]
    (cX, cY) = (w//2, h//2)
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    nW = int((h*sin)+(w*cos))
    nH = int((h*cos)+(w*sin))
    M[0, 2] += (nW/2) - cX
    M[1, 2] += (nH/2) - cY
    return cv2.warpAffine(image, M, (nW, nH))
    # return cv2.warpAffine(image, M, (h, w))

def rotate_bound_1(image, angle):
    '''
    旋转后图像尺寸不变，图像不一定都能展现
    :param image: 原图
    :param angle: 角度
    :return: 旋转后的图像
    '''
    (h, w) = image.shape[:2]
    (cX, cY) = (w//2, h//2)
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    # cos = np.abs(M[0, 0])
    # sin = np.abs(M[0, 1])
    # nW = int((h*sin)+(w*cos))
    # nH = int((h*cos)+(w*sin))
    # M[0, 2] += (nW/2) - cX
    # M[1, 2] += (nH/2) - cY
    # return cv2.warpAffine(image, M, (nW, nH))
    return cv2.warpAffine(image, M, (h, w))

def create_qingxie(images_path, output_path, img, out_row_visible_img, out_col_visible_img, out_row_invisible_img, out_col_invisible_img, i):
    '''
    直接旋转已生成的标签图像
    :param images_path:
    :param output_path:
    :param img:
    :param out_row_visible_img:
    :param out_col_visible_img:
    :param out_row_invisible_img:
    :param out_col_invisible_img:
    :param i:
    :return:
    '''
    angle = np.random.randint(-10, 10)
    row_qingxie = rotate_bound(out_row_visible_img, angle)
    col_qingxie = rotate_bound(out_col_visible_img, angle)
    nrow_qingxie = rotate_bound(out_row_invisible_img, angle)
    # cv2.namedWindow('nrow_qingxie', cv2.WINDOW_NORMAL)
    # cv2.imshow('nrow_qingxie', nrow_qingxie*255)
    # cv2.waitKey(0)
    # cv2.destroyWindow('nrow_qingxie')
    ncol_qingxie = rotate_bound(out_col_invisible_img, angle)
    img_qingxie = rotate_bound(img, angle)

    # out_row_qingxie_1 = rotate_bound_1(out_row_visible_img, angle)
    # out_col_qingxie_1 = rotate_bound_1(out_col_visible_img, angle)
    # out_nrow_qingxie_1 = rotate_bound_1(out_row_invisible_img, angle)
    # out_ncol_qingxie_1 = rotate_bound_1(out_col_invisible_img, angle)
    # out_img_qingxie_1 = rotate_bound_1(out_img, angle)

    # 裁剪图像到固定尺寸大小到512*512，采用cv2.INTER_NEAREST方法
    out_row_qingxie = cv2.resize(row_qingxie, (256, 256), interpolation=cv2.INTER_AREA)
    # print('行可视化图像不同值:', out_row_visible_img)
    out_col_qingxie = cv2.resize(col_qingxie, (256, 256), interpolation=cv2.INTER_AREA)
    # print('列可视化图像不同值:', np.unique(out_col_visible_img))
    out_nrow_qingxie = cv2.resize(nrow_qingxie, (256, 256), interpolation=cv2.INTER_AREA)
    # print('行不可视化图像不同值:', np.unique(out_row_invisible_img))
    out_ncol_qingxie = cv2.resize(ncol_qingxie, (256, 256), interpolation=cv2.INTER_AREA)
    # print('列不可视化图像不同值:', np.unique(out_col_invisible_img))
    out_img_qingxie = cv2.resize(img_qingxie, (512, 512), interpolation=cv2.INTER_AREA)
    # # 裁剪图像到固定尺寸大小到512*512，采用cv2.INTER_NEAREST方法
    # out_row_qingxie_1 = cv2.resize(out_row_qingxie_1, (int(512), int(512)), interpolation=cv2.INTER_NEAREST)
    # # print('行可视化图像不同值:', out_row_visible_img)
    # out_col_qingxie_1 = cv2.resize(out_col_qingxie_1, (int(512), int(512)), interpolation=cv2.INTER_NEAREST)
    # # print('列可视化图像不同值:', np.unique(out_col_visible_img))
    # out_nrow_qingxie_1 = cv2.resize(out_nrow_qingxie_1, (int(512), int(512)), interpolation=cv2.INTER_NEAREST)
    # # print('行不可视化图像不同值:', np.unique(out_row_invisible_img))
    # out_ncol_qingxie_1 = cv2.resize(out_ncol_qingxie_1, (int(512), int(512)), interpolation=cv2.INTER_NEAREST)
    # # print('列不可视化图像不同值:', np.unique(out_col_invisible_img))
    # out_img_qingxie_1 = cv2.resize(out_img_qingxie_1, (int(512), int(512)), interpolation=cv2.INTER_NEAREST)

    cv2.imwrite(output_path + 'label_row' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.png', out_row_qingxie)
    cv2.imwrite(output_path + 'label_col' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.png', out_col_qingxie)
    cv2.imwrite(output_path + 'label_nrow' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.png', out_nrow_qingxie)
    cv2.imwrite(output_path + 'label_ncol' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.png', out_ncol_qingxie)
    cv2.imwrite(output_path + 'image' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.jpg', out_img_qingxie)

def rotate_point(points_list, out_img, M):
    img = out_img.copy()
    for point in points_list:
        point_rot = []
        for x, y in point:
            nX = int(x * M[0, 0] + y * M[0, 1] + M[0, 2])
            nY = int(x * M[1, 0] + y * M[1, 1] + M[1, 2])
            point_rot.append([nX, nY])
        point_rot_array = np.array(point_rot)
        cv2.fillPoly(img, [point_rot_array], (1,))
    # out_row_visible_img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
    return img

def create_qingxie_1(images_path, output_path, img, row_visible_point, col_visible_point, row_invisible_point, col_invisible_point, i):
    '''
    先旋转四个顶点坐标，再连线，生成标签图像
    :param images_path:
    :param output_path:
    :param img:
    :param row_visible_point:
    :param col_visible_point:
    :param row_invisible_point:
    :param col_invisible_point:
    :param i:
    :return:
    '''
    angle = np.random.randint(-10, 10)
    h, w, _ = img.shape
    (cX, cY) = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D((cX, cY), angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY
    img = cv2.warpAffine(img, M, (nW, nH))
    # out_img = np.zeros((img.shape[0], img.shape[1], 1), np.uint8)
    out_img = np.zeros((img.shape[0], img.shape[1]), np.uint8)

    row_visible_img = rotate_point(row_visible_point, out_img, M)
    row_invisible_img = rotate_point(row_invisible_point, out_img, M)
    col_visible_img = rotate_point(col_visible_point, out_img, M)
    col_invisible_img = rotate_point(col_invisible_point, out_img, M)

    out_row_visible_img = Image.fromarray(row_visible_img).resize((256, 256), Image.ANTIALIAS)
    out_row_visible_img.save(output_path + 'label_row' + '/' + images_path[i].split('.')[0].split \
        ('\\')[-1] + '_qingxie.png')
    out_row_invisible_img = Image.fromarray(row_invisible_img).resize((256, 256), Image.ANTIALIAS)
    out_row_invisible_img.save(output_path + 'label_nrow' + '/' + images_path[i].split('.')[0].split \
        ('\\')[-1] + '_qingxie.png')
    out_col_visible_img = Image.fromarray(col_visible_img).resize((256, 256), Image.ANTIALIAS)
    out_col_visible_img.save(output_path + 'label_col' + '/' + images_path[i].split('.')[0].split \
        ('\\')[-1] + '_qingxie.png')
    out_col_invisible_img = Image.fromarray(col_invisible_img).resize((256, 256), Image.ANTIALIAS)
    out_col_invisible_img.save(output_path + 'label_ncol' + '/' + images_path[i].split('.')[0].split \
        ('\\')[-1] + '_qingxie.png')
    out_img = Image.fromarray(img[:, :, ::-1]).resize((512, 512), Image.ANTIALIAS)
    out_img.save(output_path + 'image' + '/' + images_path[i].split('.')[0].split \
        ('\\')[-1] + '_qingxie.jpg')

    # out_row_visible_img = cv2.resize(row_visible_img, (256, 256), interpolation=cv2.INTER_AREA)
    # out_col_visible_img = cv2.resize(col_visible_img, (256, 256), interpolation=cv2.INTER_AREA)
    # out_row_invisible_img = cv2.resize(row_invisible_img, (256, 256), interpolation=cv2.INTER_AREA)
    # out_col_invisible_img = cv2.resize(col_invisible_img, (256, 256), interpolation=cv2.INTER_AREA)
    # out_img = cv2.resize(img, (int(512), int(512)), interpolation=cv2.INTER_AREA)

    # cv2.imwrite(output_path + 'label_row' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.png',
    #             out_row_visible_img)
    # cv2.imwrite(output_path + 'label_col' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.png',
    #             out_col_visible_img)
    # cv2.imwrite(output_path + 'label_nrow' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.png',
    #             out_row_invisible_img*255)
    # cv2.imwrite(output_path + 'label_ncol' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.png',
    #             out_col_invisible_img)
    # cv2.imwrite(output_path + 'image' + '/' + images_path[i].split('.')[0].split('\\')[-1] + '_qingxie.jpg', out_img)

def main(base_path, output_path, augment_qingxie=True):
    # 创建输出文件路径
    subdirs = ['label_row', 'label_nrow', 'label_col', 'label_ncol', 'image']
    [os.mkdir(output_path+subdir) for subdir in subdirs if not os.path.exists(output_path+subdir)]

    # 得到图像和json文件名列表
    images_path = glob.glob(base_path+'*.jpg')
    jsons_path = [image_path.replace('jpg', 'json') for image_path in images_path if os.path.exists(image_path.replace('jpg', 'json'))]

    # 定义存储四种框线的列表变量
    row_visible = list()
    col_visible = list()
    row_invisible = list()
    col_invisible = list()

    # 遍历打标签的json文件
    for i in range(len(jsons_path)):
        # 清空四种框线列表变量信息
        row_visible.clear()
        col_visible.clear()
        row_invisible.clear()
        col_invisible.clear()

        # 读取图像信息和标签信息
        img_array = cv2.imread(images_path[i])
        h, w, _ = img_array.shape
        # out_img = np.ones((h, w, 1), np.uint8)
        out_img = np.zeros((h, w, 1), np.uint8)
        with open(jsons_path[i]) as f:
            shapes = json.load(f)['shapes']
            for shape in shapes:
                label = shape['label']
                points = shape['points']

                # 对每个标注框进行标签的判断
                if label == 'row_visible':
                    row_visible.append(np.array(points, dtype=np.int))
                elif label == 'col_visible':
                    col_visible.append(np.array(points, dtype=np.int))
                elif label == 'row_invisible':
                    row_invisible.append(np.array(points, dtype=np.int))
                else:
                    col_invisible.append(np.array(points, dtype=np.int))

        # 在copy的输出图像上画多边形, fillPoly可以一次填充多个图形cv2.fillPoly(img, [area1, area2], (255, 255, 255))，当前填充值为1
        row_visible_img = cv2.fillPoly(out_img.copy(), row_visible, color=(1, ))
        col_visible_img = cv2.fillPoly(out_img.copy(), col_visible, color=(1, ))
        row_invisible_img = cv2.fillPoly(out_img.copy(), row_invisible, color=(1, ))
        col_invisible_img = cv2.fillPoly(out_img.copy(), col_invisible, color=(1, ))

        # 裁剪图像到固定尺寸大小到512*512，采用cv2.INTER_NEAREST方法
        out_row_visible_img = cv2.resize(row_visible_img, (256, 256), interpolation=cv2.INTER_AREA)
        # print('行可视化图像不同值:', out_row_visible_img)
        out_col_visible_img = cv2.resize(col_visible_img, (256, 256), interpolation=cv2.INTER_AREA)
        # print('列可视化图像不同值:', np.unique(out_col_visible_img))
        out_row_invisible_img = cv2.resize(row_invisible_img, (256, 256), interpolation=cv2.INTER_AREA)
        # cv2.namedWindow('nrow', cv2.WINDOW_NORMAL)
        # cv2.imshow('nrow', row_invisible_img*255)
        # cv2.waitKey(0)
        # cv2.destroyWindow('nrow')
        # print('行不可视化图像不同值:', np.unique(out_row_invisible_img))
        out_col_invisible_img = cv2.resize(col_invisible_img, (256, 256), interpolation=cv2.INTER_AREA)
        # print('列不可视化图像不同值:', np.unique(out_col_invisible_img))
        out_img = cv2.resize(img_array, (int(512), int(512)), interpolation=cv2.INTER_AREA)

        # 将输出图像分别保存到输出目录下，原始图像为jpg，标签图像为png
        cv2.imwrite(output_path+'label_row'+'/'+images_path[i].split('.')[0].split('\\')[-1]+'.png', out_row_visible_img)
        cv2.imwrite(output_path+'label_col'+'/'+images_path[i].split('.')[0].split('\\')[-1]+'.png', out_col_visible_img)
        cv2.imwrite(output_path+'label_nrow'+'/'+images_path[i].split('.')[0].split('\\')[-1]+'.png', out_row_invisible_img)
        cv2.imwrite(output_path+'label_ncol'+'/'+images_path[i].split('.')[0].split('\\')[-1]+'.png', out_col_invisible_img)
        cv2.imwrite(output_path+'image'+'/'+images_path[i].split('.')[0].split('\\')[-1]+'.jpg', out_img)

        # Add 2021/10/27 对数据进行增强操作（小角度倾斜）
        if augment_qingxie:
            # create_qingxie(images_path, output_path, out_img, out_row_visible_img, out_col_visible_img, out_row_invisible_img, out_col_invisible_img, i)
            create_qingxie_1(images_path, output_path, img_array, row_visible, col_visible, row_invisible, col_invisible, i)

if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])
